{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "58f3e006-b933-44f6-8628-b241794a8a26",
   "metadata": {},
   "source": [
    "# PEFT 库 QLoRA 实战 - ChatGLM3-6B  使用10K数据训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e8e39e0d-6c4f-4ba1-a9b6-11c1f179f1c1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:16.911870Z",
     "iopub.status.busy": "2025-08-08T06:13:16.911618Z",
     "iopub.status.idle": "2025-08-08T06:13:16.919510Z",
     "shell.execute_reply": "2025-08-08T06:13:16.918964Z"
    }
   },
   "outputs": [],
   "source": [
    "# 定义全局变量和参数\n",
    "model_name_or_path = 'THUDM/chatglm3-6b'  # 模型ID或本地路径\n",
    "train_data_path = 'HasturOfficial/adgen'    # 训练数据路径\n",
    "eval_data_path = None                     # 验证数据路径，如果没有则设置为None\n",
    "seed = 8                                 # 随机种子\n",
    "max_input_length = 512                    # 输入的最大长度\n",
    "max_output_length = 1536                  # 输出的最大长度\n",
    "lora_rank = 4                             # LoRA秩\n",
    "lora_alpha = 32                           # LoRA alpha值\n",
    "lora_dropout = 0.05                       # LoRA Dropout率\n",
    "resume_from_checkpoint = None             # 如果从checkpoint恢复训练，指定路径\n",
    "prompt_text = ''                          # 所有数据前的指令文本\n",
    "compute_dtype = 'fp32'                    # 计算数据类型（fp32, fp16, bf16）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "207f68a3-cd38-4530-b391-77957762b435",
   "metadata": {},
   "source": [
    "# 数据准备"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e4cf1c6f-77f4-44e5-af05-9f5d1d61392c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:16.921559Z",
     "iopub.status.busy": "2025-08-08T06:13:16.921335Z",
     "iopub.status.idle": "2025-08-08T06:13:20.410904Z",
     "shell.execute_reply": "2025-08-08T06:13:20.409958Z"
    }
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(train_data_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "08614e44-8cbe-4730-a9a9-582149a38684",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:20.415719Z",
     "iopub.status.busy": "2025-08-08T06:13:20.414190Z",
     "iopub.status.idle": "2025-08-08T06:13:20.426388Z",
     "shell.execute_reply": "2025-08-08T06:13:20.425645Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['content', 'summary'],\n",
       "        num_rows: 114599\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['content', 'summary'],\n",
       "        num_rows: 1070\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "adb7e045-fd80-4ef9-9f8c-2a8daf262cd4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:20.430761Z",
     "iopub.status.busy": "2025-08-08T06:13:20.429572Z",
     "iopub.status.idle": "2025-08-08T06:13:20.438863Z",
     "shell.execute_reply": "2025-08-08T06:13:20.438313Z"
    }
   },
   "outputs": [],
   "source": [
    "from datasets import ClassLabel, Sequence\n",
    "import random\n",
    "import pandas as pd\n",
    "from IPython.display import display, HTML\n",
    "\n",
    "def show_random_elements(dataset, num_examples=10):\n",
    "    assert num_examples <= len(dataset), \"Can't pick more elements than there are in the dataset.\"\n",
    "    picks = []\n",
    "    for _ in range(num_examples):\n",
    "        pick = random.randint(0, len(dataset)-1)\n",
    "        while pick in picks:\n",
    "            pick = random.randint(0, len(dataset)-1)\n",
    "        picks.append(pick)\n",
    "    \n",
    "    df = pd.DataFrame(dataset[picks])\n",
    "    for column, typ in dataset.features.items():\n",
    "        if isinstance(typ, ClassLabel):\n",
    "            df[column] = df[column].transform(lambda i: typ.names[i])\n",
    "        elif isinstance(typ, Sequence) and isinstance(typ.feature, ClassLabel):\n",
    "            df[column] = df[column].transform(lambda x: [typ.feature.names[i] for i in x])\n",
    "    display(HTML(df.to_html()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cea7bc98-775b-4592-afed-4ce7d3f5153b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:20.442030Z",
     "iopub.status.busy": "2025-08-08T06:13:20.441187Z",
     "iopub.status.idle": "2025-08-08T06:13:20.449072Z",
     "shell.execute_reply": "2025-08-08T06:13:20.448530Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>content</th>\n",
       "      <th>summary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>类型#裤*材质#网纱*材质#蕾丝*颜色#红色*颜色#裸色*风格#性感*图案#蕾丝*裤款式#纱网</td>\n",
       "      <td>正红色的蕾丝蔓延扩散在裸色的网纱上，犹如蔷薇&lt;UNK&gt;下开启的一场&lt;UNK&gt;。搭配挂脖式设计，增强整体的稳定性，同时也有着呼之欲出的妩媚感。配套的t裤以正红色蕾丝遮掩私密，裸色网纱犹如裙摆围绕。再有缎面感的面料制成t&lt;UNK&gt;档，柔软贴身也性感撩人。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>类型#上衣*版型#立体剪裁*颜色#黑色*颜色#军绿色*图案#字母*图案#文字*图案#印花*衣样式#风衣*衣款式#抽绳</td>\n",
       "      <td>来自BRAND品牌的这款男士风衣，黑色、军绿色两款颜色，立体剪裁，彰显男士气质。后背的字母印花，彰显品牌辨识度，腰间抽绳，方便调节，更利落有型。</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>类型#裙*版型#宽松*材质#棉*图案#条纹*图案#印花*裙型#直筒裙*裙袖型#喇叭袖</td>\n",
       "      <td>BRAND&lt;UNK&gt;，不能少的事情当然是把BRAND的大头像印在睡裙上。BRAND的印花&lt;UNK&gt;&lt;UNK&gt;的样子特别可爱，就像小女孩&lt;UNK&gt;&lt;UNK&gt;故事一样。这款睡裙采用直筒宽松的版型，纯棉的质地提高了穿着体验。喇叭袖的设计更显手臂纤细，条纹纹理的安排不会让整体看上去显得单调。是一款满足少女心的睡衣。</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_random_elements(dataset[\"train\"], num_examples=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4694412-bc49-4030-a770-08b6932a86c2",
   "metadata": {},
   "source": [
    "# 使用 ChatGLM3-6b Tokenizer 处理数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "de2c8d4c-5a6b-481d-913b-23d2df3c956e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:20.451196Z",
     "iopub.status.busy": "2025-08-08T06:13:20.450967Z",
     "iopub.status.idle": "2025-08-08T06:13:23.093155Z",
     "shell.execute_reply": "2025-08-08T06:13:23.092408Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/huggingface_hub/file_download.py:945: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "# revision='b098244' 版本对应的 ChatGLM3-6B 设置 use_reentrant=False\n",
    "# 最新版本 use_reentrant 被设置为 True，会增加不必要的显存开销\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,\n",
    "                                          trust_remote_code=True,\n",
    "                                          revision='b098244')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3209ad6d-fe52-4afe-8158-e0a565c691a1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:23.096013Z",
     "iopub.status.busy": "2025-08-08T06:13:23.095734Z",
     "iopub.status.idle": "2025-08-08T06:13:23.102326Z",
     "shell.execute_reply": "2025-08-08T06:13:23.101735Z"
    }
   },
   "outputs": [],
   "source": [
    "# tokenize_func 函数\n",
    "def tokenize_func(example, tokenizer, ignore_label_id=-100):\n",
    "    \"\"\"\n",
    "    对单个数据样本进行tokenize处理。\n",
    "\n",
    "    参数:\n",
    "    example (dict): 包含'content'和'summary'键的字典，代表训练数据的一个样本。\n",
    "    tokenizer (transformers.PreTrainedTokenizer): 用于tokenize文本的tokenizer。\n",
    "    ignore_label_id (int, optional): 在label中用于填充的忽略ID，默认为-100。\n",
    "\n",
    "    返回:\n",
    "    dict: 包含'tokenized_input_ids'和'labels'的字典，用于模型训练。\n",
    "    \"\"\"\n",
    "\n",
    "    # 构建问题文本\n",
    "    question = prompt_text + example['content']\n",
    "    if example.get('input', None) and example['input'].strip():\n",
    "        question += f'\\n{example[\"input\"]}'\n",
    "\n",
    "    # 构建答案文本\n",
    "    answer = example['summary']\n",
    "\n",
    "    # 对问题和答案文本进行tokenize处理\n",
    "    q_ids = tokenizer.encode(text=question, add_special_tokens=False)\n",
    "    a_ids = tokenizer.encode(text=answer, add_special_tokens=False)\n",
    "\n",
    "    # 如果tokenize后的长度超过最大长度限制，则进行截断\n",
    "    if len(q_ids) > max_input_length - 2:  # 保留空间给gmask和bos标记\n",
    "        q_ids = q_ids[:max_input_length - 2]\n",
    "    if len(a_ids) > max_output_length - 1:  # 保留空间给eos标记\n",
    "        a_ids = a_ids[:max_output_length - 1]\n",
    "\n",
    "    # 构建模型的输入格式\n",
    "    input_ids = tokenizer.build_inputs_with_special_tokens(q_ids, a_ids)\n",
    "    question_length = len(q_ids) + 2  # 加上gmask和bos标记\n",
    "\n",
    "    # 构建标签，对于问题部分的输入使用ignore_label_id进行填充\n",
    "    labels = [ignore_label_id] * question_length + input_ids[question_length:]\n",
    "\n",
    "    return {'input_ids': input_ids, 'labels': labels}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6f190bdc-6257-4d94-9eee-10eda43ce44e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:23.104313Z",
     "iopub.status.busy": "2025-08-08T06:13:23.104088Z",
     "iopub.status.idle": "2025-08-08T06:13:23.122116Z",
     "shell.execute_reply": "2025-08-08T06:13:23.121505Z"
    }
   },
   "outputs": [],
   "source": [
    "column_names = dataset['train'].column_names\n",
    "tokenized_dataset = dataset['train'].map(\n",
    "    lambda example: tokenize_func(example, tokenizer),\n",
    "    batched=False, \n",
    "    remove_columns=column_names\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3d97058b-5a92-433b-b3e6-5085d941e732",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:23.124220Z",
     "iopub.status.busy": "2025-08-08T06:13:23.123982Z",
     "iopub.status.idle": "2025-08-08T06:13:23.130888Z",
     "shell.execute_reply": "2025-08-08T06:13:23.130347Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>input_ids</th>\n",
       "      <th>labels</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[64790, 64792, 30910, 33467, 31010, 49534, 30998, 32799, 31010, 40512, 30998, 37505, 31010, 54827, 54583, 30998, 55500, 46025, 31010, 38712, 54764, 30998, 55500, 46025, 31010, 42373, 30998, 55500, 54811, 58709, 31010, 55097, 55759, 30998, 55500, 40877, 31010, 56069, 54762, 30998, 55500, 40877, 31010, 42875, 30998, 55500, 40877, 31010, 55097, 55759, 30910, 54713, 55115, 42373, 32195, 38712, 54764, 42064, 48233, 56597, 55857, 32755, 32698, 54827, 32441, 46025, 31123, 54783, 42373, 56310, 55297, 54715, 54888, 47760, 40512, 54963, 55400, 33253, 54656, 31794, 42373, 34119, 54706, 54783, 42373, 47760, 34372, 42524, 31155, 32413, 31123, 54713, 55115, 42373, 54656, 32195, 55097, 55759, 54811, 58709, ...]</td>\n",
       "      <td>[-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 30910, 54713, 55115, 42373, 32195, 38712, 54764, 42064, 48233, 56597, 55857, 32755, 32698, 54827, 32441, 46025, 31123, 54783, 42373, 56310, 55297, 54715, 54888, 47760, 40512, 54963, 55400, 33253, 54656, 31794, 42373, 34119, 54706, 54783, 42373, 47760, 34372, 42524, 31155, 32413, 31123, 54713, 55115, 42373, 54656, 32195, 55097, 55759, 54811, 58709, ...]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_random_elements(tokenized_dataset, num_examples=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a31da86-c865-4819-a4a5-8edd3b76399c",
   "metadata": {},
   "source": [
    "# 数据集处理：shuffle & flatten"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f2c1fd72-c565-49bf-8462-27c1595e96f1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:23.132873Z",
     "iopub.status.busy": "2025-08-08T06:13:23.132636Z",
     "iopub.status.idle": "2025-08-08T06:13:23.137766Z",
     "shell.execute_reply": "2025-08-08T06:13:23.137218Z"
    }
   },
   "outputs": [],
   "source": [
    "tokenized_dataset = tokenized_dataset.shuffle(seed=seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "51a71cfd-1edb-4883-bd4c-677bf317d826",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:23.139673Z",
     "iopub.status.busy": "2025-08-08T06:13:23.139451Z",
     "iopub.status.idle": "2025-08-08T06:13:23.146861Z",
     "shell.execute_reply": "2025-08-08T06:13:23.146313Z"
    }
   },
   "outputs": [],
   "source": [
    "tokenized_dataset = tokenized_dataset.flatten_indices()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8837656-4454-4dc9-be63-94e3efdc44ec",
   "metadata": {},
   "source": [
    "# 定义 DataCollatorForChatGLM 类 批量处理数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "72feeed0-5ee1-438d-b619-a4c027b91034",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:23.149009Z",
     "iopub.status.busy": "2025-08-08T06:13:23.148780Z",
     "iopub.status.idle": "2025-08-08T06:13:23.156143Z",
     "shell.execute_reply": "2025-08-08T06:13:23.155613Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from typing import List, Dict, Optional\n",
    "\n",
    "# DataCollatorForChatGLM 类\n",
    "class DataCollatorForChatGLM:\n",
    "    \"\"\"\n",
    "    用于处理批量数据的DataCollator，尤其是在使用 ChatGLM 模型时。\n",
    "\n",
    "    该类负责将多个数据样本（tokenized input）合并为一个批量，并在必要时进行填充(padding)。\n",
    "\n",
    "    属性:\n",
    "    pad_token_id (int): 用于填充(padding)的token ID。\n",
    "    max_length (int): 单个批量数据的最大长度限制。\n",
    "    ignore_label_id (int): 在标签中用于填充的ID。\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, pad_token_id: int, max_length: int = 2048, ignore_label_id: int = -100):\n",
    "        \"\"\"\n",
    "        初始化DataCollator。\n",
    "\n",
    "        参数:\n",
    "        pad_token_id (int): 用于填充(padding)的token ID。\n",
    "        max_length (int): 单个批量数据的最大长度限制。\n",
    "        ignore_label_id (int): 在标签中用于填充的ID，默认为-100。\n",
    "        \"\"\"\n",
    "        self.pad_token_id = pad_token_id\n",
    "        self.ignore_label_id = ignore_label_id\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __call__(self, batch_data: List[Dict[str, List]]) -> Dict[str, torch.Tensor]:\n",
    "        \"\"\"\n",
    "        处理批量数据。\n",
    "\n",
    "        参数:\n",
    "        batch_data (List[Dict[str, List]]): 包含多个样本的字典列表。\n",
    "\n",
    "        返回:\n",
    "        Dict[str, torch.Tensor]: 包含处理后的批量数据的字典。\n",
    "        \"\"\"\n",
    "        # 计算批量中每个样本的长度\n",
    "        len_list = [len(d['input_ids']) for d in batch_data]\n",
    "        batch_max_len = max(len_list)  # 找到最长的样本长度\n",
    "\n",
    "        input_ids, labels = [], []\n",
    "        for len_of_d, d in sorted(zip(len_list, batch_data), key=lambda x: -x[0]):\n",
    "            pad_len = batch_max_len - len_of_d  # 计算需要填充的长度\n",
    "            # 添加填充，并确保数据长度不超过最大长度限制\n",
    "            ids = d['input_ids'] + [self.pad_token_id] * pad_len\n",
    "            label = d['labels'] + [self.ignore_label_id] * pad_len\n",
    "            if batch_max_len > self.max_length:\n",
    "                ids = ids[:self.max_length]\n",
    "                label = label[:self.max_length]\n",
    "            input_ids.append(torch.LongTensor(ids))\n",
    "            labels.append(torch.LongTensor(label))\n",
    "\n",
    "        # 将处理后的数据堆叠成一个tensor\n",
    "        input_ids = torch.stack(input_ids)\n",
    "        labels = torch.stack(labels)\n",
    "\n",
    "        return {'input_ids': input_ids, 'labels': labels}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c5428132-8cb1-48de-9076-650418cce968",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:23.158022Z",
     "iopub.status.busy": "2025-08-08T06:13:23.157801Z",
     "iopub.status.idle": "2025-08-08T06:13:23.160445Z",
     "shell.execute_reply": "2025-08-08T06:13:23.159921Z"
    }
   },
   "outputs": [],
   "source": [
    "# 准备数据整理器\n",
    "data_collator = DataCollatorForChatGLM(pad_token_id=tokenizer.pad_token_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5d00321-5935-4f65-a4f9-9700c15c2356",
   "metadata": {},
   "source": [
    "# 训练模型\n",
    "## 加载 ChatGLM3-6B 量化模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f040725d-5e9d-4c7b-a70b-28fef78250c2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:23.162373Z",
     "iopub.status.busy": "2025-08-08T06:13:23.162135Z",
     "iopub.status.idle": "2025-08-08T06:13:23.179237Z",
     "shell.execute_reply": "2025-08-08T06:13:23.178647Z"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModel, BitsAndBytesConfig\n",
    "\n",
    "_compute_dtype_map = {\n",
    "    'fp32': torch.float32,\n",
    "    'fp16': torch.float16,\n",
    "    'bf16': torch.bfloat16\n",
    "}\n",
    "\n",
    "# QLoRA 量化配置\n",
    "q_config = BitsAndBytesConfig(load_in_4bit=True,\n",
    "                              bnb_4bit_quant_type='nf4',\n",
    "                              bnb_4bit_use_double_quant=True,\n",
    "                              bnb_4bit_compute_dtype=_compute_dtype_map['bf16'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "300d5d0e-da9e-4a88-9762-07de9f694369",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:23.181428Z",
     "iopub.status.busy": "2025-08-08T06:13:23.181196Z",
     "iopub.status.idle": "2025-08-08T06:13:36.957139Z",
     "shell.execute_reply": "2025-08-08T06:13:36.956417Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/huggingface_hub/file_download.py:945: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b9d04fe8e4d64dce8d0ce6481e3a83b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/huggingface_hub/file_download.py:945: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# revision='b098244' 版本对应的 ChatGLM3-6B 设置 use_reentrant=False\n",
    "# 最新版本 use_reentrant 被设置为 True，会增加不必要的显存开销\n",
    "model = AutoModel.from_pretrained(model_name_or_path,\n",
    "                                  quantization_config=q_config,\n",
    "                                  device_map='auto',\n",
    "                                  trust_remote_code=True,\n",
    "                                  revision='b098244')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3519e35f-11e1-4163-8cbf-7c8eaa7a4a83",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:36.960161Z",
     "iopub.status.busy": "2025-08-08T06:13:36.959697Z",
     "iopub.status.idle": "2025-08-08T06:13:36.966022Z",
     "shell.execute_reply": "2025-08-08T06:13:36.965375Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3739.69MiB\n"
     ]
    }
   ],
   "source": [
    "# 获取当前模型占用的 GPU显存（差值为预留给 PyTorch 的显存）\n",
    "memory_footprint_bytes = model.get_memory_footprint()\n",
    "memory_footprint_mib = memory_footprint_bytes / (1024 ** 2)  # 转换为 MiB\n",
    "\n",
    "print(f\"{memory_footprint_mib:.2f}MiB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e3e53c8-4ad2-489d-ace2-4d0ccbdd5649",
   "metadata": {},
   "source": [
    "# 预处理量化模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c1b90c77-2ac2-478a-b8e5-5df5a8662396",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:36.968134Z",
     "iopub.status.busy": "2025-08-08T06:13:36.967904Z",
     "iopub.status.idle": "2025-08-08T06:13:37.021452Z",
     "shell.execute_reply": "2025-08-08T06:13:37.020711Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using an old version of the checkpointing format that is deprecated (We will also silently ignore `gradient_checkpointing_kwargs` in case you passed it).Please update to the new format on your modeling file. To use the new format, you need to completely remove the definition of the method `_set_gradient_checkpointing` in your model.\n"
     ]
    }
   ],
   "source": [
    "from peft import TaskType, LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "\n",
    "kbit_model = prepare_model_for_kbit_training(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "77f1ac47-44be-4489-a0b0-3369fd5c05bb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:37.023733Z",
     "iopub.status.busy": "2025-08-08T06:13:37.023479Z",
     "iopub.status.idle": "2025-08-08T06:13:37.026789Z",
     "shell.execute_reply": "2025-08-08T06:13:37.026175Z"
    }
   },
   "outputs": [],
   "source": [
    "from peft.utils import TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING\n",
    "\n",
    "target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING['chatglm']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "76de3bf1-58bc-4970-813e-cab75f9587d4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:37.028822Z",
     "iopub.status.busy": "2025-08-08T06:13:37.028598Z",
     "iopub.status.idle": "2025-08-08T06:13:37.032439Z",
     "shell.execute_reply": "2025-08-08T06:13:37.031915Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['query_key_value']"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_modules"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c69ca723-2692-4fb4-a398-203c64fa02d3",
   "metadata": {},
   "source": [
    "# LoRA  适配器配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "98533ed6-5cd1-4cdd-a3ae-78b77917b8f7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:37.034409Z",
     "iopub.status.busy": "2025-08-08T06:13:37.034178Z",
     "iopub.status.idle": "2025-08-08T06:13:37.037519Z",
     "shell.execute_reply": "2025-08-08T06:13:37.036971Z"
    }
   },
   "outputs": [],
   "source": [
    "lora_config = LoraConfig(\n",
    "    target_modules=target_modules,\n",
    "    r=lora_rank,\n",
    "    lora_alpha=lora_alpha,\n",
    "    lora_dropout=lora_dropout,\n",
    "    bias='none',\n",
    "    inference_mode=False,\n",
    "    task_type=TaskType.CAUSAL_LM\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "fd462152-0bb6-4ba4-a742-2c39e0f3f5dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:37.039357Z",
     "iopub.status.busy": "2025-08-08T06:13:37.039123Z",
     "iopub.status.idle": "2025-08-08T06:13:37.082918Z",
     "shell.execute_reply": "2025-08-08T06:13:37.082227Z"
    }
   },
   "outputs": [],
   "source": [
    "qlora_model = get_peft_model(kbit_model, lora_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "f7f9ff60-1ff0-401d-8ade-fc4596f114c2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:37.085354Z",
     "iopub.status.busy": "2025-08-08T06:13:37.085071Z",
     "iopub.status.idle": "2025-08-08T06:13:37.090422Z",
     "shell.execute_reply": "2025-08-08T06:13:37.089824Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 974,848 || all params: 6,244,558,848 || trainable%: 0.01561115883009451\n"
     ]
    }
   ],
   "source": [
    "qlora_model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58c51e82-c560-4efd-b776-24e33e8d6f2b",
   "metadata": {},
   "source": [
    "### 训练超参数配置\n",
    "\n",
    "- 1个epoch表示对训练集的所有样本进行一次完整的训练。\n",
    "- `num_train_epochs` 表示要完整进行多少个 epochs 的训练。\n",
    "\n",
    "#### 关于使用 num_train_epochs 时，训练总步数 `steps` 的计算方法\n",
    "\n",
    "- 训练总步数： `total_steps = steps/epoch * num_train_epochs` \n",
    "- 每个epoch的训练步数：`steps/epoch = num_train_examples / (batch_size * gradient_accumulation_steps)`\n",
    "\n",
    "\n",
    "**以 `adgen` 数据集为例计算**\n",
    "\n",
    "```json\n",
    "DatasetDict({\n",
    "    train: Dataset({\n",
    "        features: ['content', 'summary'],\n",
    "        num_rows: 114599\n",
    "    })\n",
    "    validation: Dataset({\n",
    "        features: ['content', 'summary'],\n",
    "        num_rows: 1070\n",
    "    })\n",
    "})\n",
    "```\n",
    "\n",
    "代入超参数和配置进行计算：\n",
    "\n",
    "```python\n",
    "num_train_epochs = 1\n",
    "num_train_examples = 114599\n",
    "batch_size = 16\n",
    "gradient_accumulation_steps = 4\n",
    "\n",
    "\n",
    "steps = num_train_epochs * num_train_examples / (batch_size * gradient_accumulation_steps)\n",
    "      = 1 * 114599 / (16 * 4)\n",
    "      = 1790\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3915cd90-b1b5-426d-a73f-5b9421f85a1f",
   "metadata": {},
   "source": [
    "# 云服务器资源比较紧张，跑完10Kexample需要12小时，设置了只跑1K个example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "327ef16f-1c21-4788-ab9c-96156cfd1771",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:37.092587Z",
     "iopub.status.busy": "2025-08-08T06:13:37.092345Z",
     "iopub.status.idle": "2025-08-08T06:13:37.180272Z",
     "shell.execute_reply": "2025-08-08T06:13:37.179568Z"
    }
   },
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, Trainer\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=f\"models/{model_name_or_path}\",          # 输出目录\n",
    "    per_device_train_batch_size=16,                     # 每个设备的训练批量大小\n",
    "    gradient_accumulation_steps=4,                     # 梯度累积步数\n",
    "    # per_device_eval_batch_size=8,                      # 每个设备的评估批量大小\n",
    "    learning_rate=1e-3,                                # 学习率\n",
    "    #num_train_epochs=1,                                # 训练轮数\n",
    "    max_steps=1570,                 #设置10K样本的步数 1*100000/(16*4)=1563\n",
    "    lr_scheduler_type=\"linear\",                        # 学习率调度器类型\n",
    "    warmup_ratio=0.1,                                  # 预热比例\n",
    "    logging_steps=100,                                 # 日志记录步数\n",
    "    save_strategy=\"steps\",                             # 模型保存策略\n",
    "    save_steps=100,                                    # 模型保存步数\n",
    "    # evaluation_strategy=\"steps\",                       # 评估策略\n",
    "    # eval_steps=500,                                    # 评估步数\n",
    "    optim=\"adamw_torch\",                               # 优化器类型\n",
    "    fp16=True,                                        # 是否使用混合精度训练\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "44d6dbe4-853c-479a-8e06-1cb5ebee5bdf",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:37.182959Z",
     "iopub.status.busy": "2025-08-08T06:13:37.182491Z",
     "iopub.status.idle": "2025-08-08T06:13:37.191181Z",
     "shell.execute_reply": "2025-08-08T06:13:37.190594Z"
    }
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "        model=qlora_model,\n",
    "        args=training_args,\n",
    "        train_dataset=tokenized_dataset,\n",
    "        data_collator=data_collator\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "475f1768-cd23-47b9-aa04-d04f35fa6149",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T06:13:37.193257Z",
     "iopub.status.busy": "2025-08-08T06:13:37.193012Z",
     "iopub.status.idle": "2025-08-08T18:48:39.242307Z",
     "shell.execute_reply": "2025-08-08T18:48:39.241638Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1570' max='1570' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1570/1570 12:34:36, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>3.901400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>3.330300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>3.239900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>3.213400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>500</td>\n",
       "      <td>3.167100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>3.150700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>700</td>\n",
       "      <td>3.143500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>3.119200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>900</td>\n",
       "      <td>3.104300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>3.075200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1100</td>\n",
       "      <td>3.060700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1200</td>\n",
       "      <td>3.067100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1300</td>\n",
       "      <td>3.049900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>3.045000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1500</td>\n",
       "      <td>3.027600</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/lm_ai_learn/lib/python3.10/site-packages/torch/utils/checkpoint.py:464: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=1570, training_loss=3.173244500615794, metrics={'train_runtime': 45301.7988, 'train_samples_per_second': 2.218, 'train_steps_per_second': 0.035, 'total_flos': 6.067844113365074e+17, 'train_loss': 3.173244500615794, 'epoch': 0.88})"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "2dfa0cfc-ff6b-4b19-beb3-c3bb91d18537",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-08-08T18:48:39.244954Z",
     "iopub.status.busy": "2025-08-08T18:48:39.244490Z",
     "iopub.status.idle": "2025-08-08T18:48:39.290670Z",
     "shell.execute_reply": "2025-08-08T18:48:39.289974Z"
    }
   },
   "outputs": [],
   "source": [
    "trainer.model.save_pretrained(f\"models/demo/{model_name_or_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "399db311-ee20-483f-8bff-fb696efcad3d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {
     "1a81f32ff34945ca9e52110d0f822bfb": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "FloatProgressModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "FloatProgressModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "ProgressView",
       "bar_style": "success",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_ed4f939fc1534e21bfb5905120ed52b1",
       "max": 7.0,
       "min": 0.0,
       "orientation": "horizontal",
       "style": "IPY_MODEL_9c7ad7bc418d42a98055d607037eea28",
       "tabbable": null,
       "tooltip": null,
       "value": 7.0
      }
     },
     "3d34eb4090b545f5a686291f6269b6f2": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "5af89b6e7d3a4fcca34bc7bb521b8585": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_a68a48234c46412a8ca6b24aa1bb075c",
       "placeholder": "​",
       "style": "IPY_MODEL_a91a115cd6134fd893ad3bd72af98372",
       "tabbable": null,
       "tooltip": null,
       "value": "Loading checkpoint shards: 100%"
      }
     },
     "87cd0a8c5e5a4e70a92a03432aec4b42": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "9c7ad7bc418d42a98055d607037eea28": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "ProgressStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "ProgressStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "bar_color": null,
       "description_width": ""
      }
     },
     "a68a48234c46412a8ca6b24aa1bb075c": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "a91a115cd6134fd893ad3bd72af98372": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLStyleModel",
      "state": {
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLStyleModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "StyleView",
       "background": null,
       "description_width": "",
       "font_size": null,
       "text_color": null
      }
     },
     "b5a73eb13a6641fcb4c476c63475ef80": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HTMLModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HTMLModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HTMLView",
       "description": "",
       "description_allow_html": false,
       "layout": "IPY_MODEL_3d34eb4090b545f5a686291f6269b6f2",
       "placeholder": "​",
       "style": "IPY_MODEL_87cd0a8c5e5a4e70a92a03432aec4b42",
       "tabbable": null,
       "tooltip": null,
       "value": " 7/7 [00:05&lt;00:00,  1.47it/s]"
      }
     },
     "b7b2763ea30449e5b6f251c5d39d43d9": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     },
     "b9d04fe8e4d64dce8d0ce6481e3a83b2": {
      "model_module": "@jupyter-widgets/controls",
      "model_module_version": "2.0.0",
      "model_name": "HBoxModel",
      "state": {
       "_dom_classes": [],
       "_model_module": "@jupyter-widgets/controls",
       "_model_module_version": "2.0.0",
       "_model_name": "HBoxModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/controls",
       "_view_module_version": "2.0.0",
       "_view_name": "HBoxView",
       "box_style": "",
       "children": [
        "IPY_MODEL_5af89b6e7d3a4fcca34bc7bb521b8585",
        "IPY_MODEL_1a81f32ff34945ca9e52110d0f822bfb",
        "IPY_MODEL_b5a73eb13a6641fcb4c476c63475ef80"
       ],
       "layout": "IPY_MODEL_b7b2763ea30449e5b6f251c5d39d43d9",
       "tabbable": null,
       "tooltip": null
      }
     },
     "ed4f939fc1534e21bfb5905120ed52b1": {
      "model_module": "@jupyter-widgets/base",
      "model_module_version": "2.0.0",
      "model_name": "LayoutModel",
      "state": {
       "_model_module": "@jupyter-widgets/base",
       "_model_module_version": "2.0.0",
       "_model_name": "LayoutModel",
       "_view_count": null,
       "_view_module": "@jupyter-widgets/base",
       "_view_module_version": "2.0.0",
       "_view_name": "LayoutView",
       "align_content": null,
       "align_items": null,
       "align_self": null,
       "border_bottom": null,
       "border_left": null,
       "border_right": null,
       "border_top": null,
       "bottom": null,
       "display": null,
       "flex": null,
       "flex_flow": null,
       "grid_area": null,
       "grid_auto_columns": null,
       "grid_auto_flow": null,
       "grid_auto_rows": null,
       "grid_column": null,
       "grid_gap": null,
       "grid_row": null,
       "grid_template_areas": null,
       "grid_template_columns": null,
       "grid_template_rows": null,
       "height": null,
       "justify_content": null,
       "justify_items": null,
       "left": null,
       "margin": null,
       "max_height": null,
       "max_width": null,
       "min_height": null,
       "min_width": null,
       "object_fit": null,
       "object_position": null,
       "order": null,
       "overflow": null,
       "padding": null,
       "right": null,
       "top": null,
       "visibility": null,
       "width": null
      }
     }
    },
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
